#include "BushCalculator.h"
#include <iostream>
#include "Utility/Algebra/Vector3D.h"
#include "Utility/Algebra/Point.h"
#include "Utility/Algebra/MathTools.h"
#include "DataStructure/Common/Id.h"
#include "DataStructure/Input/Solver/DOFDirections.h"
#include "DBManager/IService/Input/Node/INodeService.h"
#include "DBManager/IService/Input/Element/IElementService.h"
#include "DBManager/IService/Input/Property/IPBUSHService.h"
#include "DBManager/IService/Result/Nodal/IDISPLACEMENTService.h"
#include "DBManager/IService/Result/ElementalForce/IBUSHForceService.h"
#include "DBManager/IService/Result/ElementalStrain/IBUSHStrainService.h"
#include "DBManager/IService/Result/ElementalStress/IBUSHStressService.h"
#include "SGFem/Pre/Assembly/Assembly.h"
#include "SGFem/Pre/Assembly/CoordinateProcessor.h"

using SG::Algebra::Matrixd;
using SG::Algebra::Vector3D;
using SG::Algebra::Point;
using SG::Algebra::PointArray;
using SG::DataStructure::Common::Status;
using SG::DataStructure::Common::Id;
using SG::DataStructure::Common::Real;
using namespace SG::DataStructure::FEM;
using SG::DataStructure::Post::BUSHForceData;
using SG::DataStructure::Post::BUSHStrainData;
using SG::DataStructure::Post::BUSHStressData;
using namespace SG::DBManager;
using SG::FEMSolver::AssemblyEigen;
using SG::FEMSolver::CoordinateProcessor;
using namespace SG::FEMSolver::Element;

void BushCalculator::Initialize (SG::DataStructure::FEM::ElementTypes type)
{
    m_type = type;
}

void BushCalculator::ComputeElementTrans
(
const std::size_t                                    nEleNodes,
const SG::DataStructure::FEM::CBUSHData&             element,
const std::vector<SG::DataStructure::FEM::NodeData>& nodeSet,
const SG::DataStructure::FEM::NodeData&              orienNode,
const SG::DataStructure::Common::Real*&              K,
_OUT bool                                            isSpecifyEleCoordSys,
_OUT SG::Algebra::Matrixd&                           TEle,
_OUT SG::Algebra::Matrixd&                           TEleOffset
)
{
    const auto& coordProcessor = CoordinateProcessor::GetInstance ();

    // 计算节点在 Basic 坐标系下的坐标
    PointArray pointSet (2);
    std::vector<std::size_t> nodeIndexSet (nEleNodes);
    for (std::size_t iNode{0}; iNode < nEleNodes; ++iNode)
    {
        Id dispCoordId;
        coordProcessor.GetGlobalCoordinate (nodeSet[iNode], _OUT pointSet[iNode], _OUT dispCoordId);
    }

    // 计算单元长度
    Vector3D xVec;
    Real L{0.0};
    if (2 == nEleNodes)
    {
        xVec = Distance (pointSet[1], pointSet[0]);
        L = std::sqrt (xVec.m_x * xVec.m_x + xVec.m_y * xVec.m_y + xVec.m_z * xVec.m_z);
    }

    isSpecifyEleCoordSys = true;
    Matrixd T (3, 3);
    const Id coordId{element.m_coordId};
    auto orien = element.GetXi ();
    if (0 == coordId)
    {
        T = SG::Algebra::identity (3);
    }
    else if (coordId > 0)
    {
        coordProcessor.GetCoordTrans (coordId, pointSet[0], _OUT T);
    }
    else if (element.m_g[1] > 0)
    {
        if (L >= 0.0001)
        {
            // 使用朝向节点
            if (element.m_go > 0)
            {
                // 计算朝向节点在 Basic 坐标系下的坐标
                Point orienPoint;
                Id dispCoordId;
                coordProcessor.GetGlobalCoordinate (orienNode, _OUT orienPoint, _OUT dispCoordId);
                // 朝向节点等价转换成 Basic 坐标系下的朝向矢量
                orien = Distance (orienPoint, pointSet[0]);
            }
            // 使用朝向矢量
            else if (!orien.IsZero ())
            {
                // 节点 A 位移输出坐标系的朝向矢量变换到 Basic 坐标系
                if (nodeSet[0].m_dispCoord > 0)
                {
                    // 获取朝向矢量的坐标系变换矩阵，即节点 A 位移输出坐标系的变换矩阵
                    Matrixd TG (3, 3);
                    coordProcessor.GetCoordTrans (nodeSet[0].m_dispCoord, pointSet[0], _OUT TG);
                    // 朝向矢量变换到 Basic 坐标系
                    orien = orien.Rotate (TG.Transpose ());
                }
            }
            else if ((0.0 == K[0] && 0.0 == K[3]) || K[1] != 0.0 || K[2] != 0.0 || K[4] != 0.0 || K[5] != 0.0)
            {
                std::cout << "Invalid BUSH element in SG::FEMSolver::Element::BushCalculator::ComputeElementMatries.\n";
                assert (0);
            }

            // 计算变换矩阵
            auto xBase{xVec};
            Real tmpL{0.0};
            xBase.Normalize (tmpL);
            // 如果单元坐标系、朝向节点、朝向矢量均未指定，则 2，3，5，6 方向的刚度系数 Ki 和阻尼系数 Bi 必须为 0.0
            if (orien.IsZero ())
            {
                isSpecifyEleCoordSys = false;
                T (0, 0) = xBase.m_x;
                T (0, 1) = xBase.m_y;
                T (0, 2) = xBase.m_z;
            }
            else
            {
                SG::Algebra::computeLocalCoordTrans (xBase, orien, _OUT T);
            }
        }
        else
        {
            std::cout << "Invalid BUSH element in SG::FEMSolver::Element::BushCalculator::ComputeElementMatries.\n";
            assert (0);
            // std::cout << "Can't support the BUSH connecting two coincident grid pointSet or one grid point SG::FEMSolver::Element::BushCalculator::ComputeElementMatries.\n";
        }
    }
    else
    {
        std::cout << "Invalid BUSH element in SG::FEMSolver::Element::BushCalculator::ComputeElementMatries.\n";
        assert (0);
    }

    // 单元节点偏置矢量
    const Id offsetCoordId{element.m_oCId};
    std::vector<Vector3D> offsetSet (2);
    if (2 == nEleNodes)
    {
        if (-1 == offsetCoordId)
        {
            if (L >= 0.0001)
            {
                const Real s{element.m_s};
                offsetSet[0] = xVec * s;
                offsetSet[1]= xVec * (s - 1);
            }
        }
        else if (offsetCoordId >= 0)
        {
            auto offset = element.GetSi ();
            if (offsetCoordId > 0)
            {
                coordProcessor.Translate2Basic (offsetCoordId, offset, _OUT offset);
            }
            for (std::size_t iNode{0}; iNode < nEleNodes; ++iNode)
            {
                offsetSet[iNode].m_x = offset.m_x - pointSet[iNode].m_x;
                offsetSet[iNode].m_y = offset.m_y - pointSet[iNode].m_y;
                offsetSet[iNode].m_z = offset.m_z - pointSet[iNode].m_z;
            }
        }
        else
        {
            std::cout << "Invalid BUSH element in SG::FEMSolver::Element::BushCalculator::ComputeElementMatries.\n";
            assert (0);
        }
    }

    // 单元变换
    for (std::size_t iNode{0}; iNode < nEleNodes; ++iNode)
    {
        const auto iNodeDof = iNode * 6;
        
        // 偏置变换矩阵
        Matrixd TOffset (3, 3);
        if (2 == nEleNodes)
        {
            TOffset (0, 1) = offsetSet[iNode].m_z;
            TOffset (2, 0) = offsetSet[iNode].m_y;
            TOffset (1, 2) = offsetSet[iNode].m_x;
            TOffset (1, 0) = -TOffset (0, 1);
            TOffset (0, 2) = -TOffset (2, 0);
            TOffset (2, 1) = -TOffset (1, 2);
        }

        // Basic 坐标系变换到位移输出坐标系
        if (0 == nodeSet[iNode].m_dispCoord)
        {
            // 单元坐标系变换到 Basic 坐标系（即位移输出坐标系）
            TEle.BlockFill (iNodeDof, iNodeDof, T);
            TEle.BlockFill (iNodeDof + 3, iNodeDof + 3, T);

            // 偏置变换
            TEleOffset.BlockFill (iNodeDof, iNodeDof, T);
            TEleOffset.BlockFill (iNodeDof + 3, iNodeDof + 3, T);
            if (2 == nEleNodes)
            {
                TEleOffset.BlockFill (iNodeDof, iNodeDof + 3, T * TOffset);
            }
        }
        else
        {
            // 获取节点位移输出坐标系坐标变换矩阵
            Matrixd TG (3, 3);
            coordProcessor.GetCoordTrans (nodeSet[iNode].m_dispCoord, pointSet[iNode], _OUT TG);
            // 单元坐标系变换到位移输出坐标系
            const auto TTemp = T * TG.Transpose ();
            TEle.BlockFill (iNodeDof, iNodeDof, TTemp);
            TEle.BlockFill (iNodeDof + 3, iNodeDof + 3, TTemp);

            // 偏置变换
            TEleOffset.BlockFill (iNodeDof, iNodeDof, TTemp);
            TEleOffset.BlockFill (iNodeDof + 3, iNodeDof + 3, TTemp);
            if (2 == nEleNodes)
            {
                TEleOffset.BlockFill (iNodeDof, iNodeDof + 3, T * TOffset * TG.Transpose ());
            }
        }
    }
}

void BushCalculator::ComputeStiffness
(
const std::size_t                        nEleNodes,
const SG::DataStructure::FEM::PBUSHData& property,
_OUT SG::Algebra::Matrixd&               kl
)
{
    Matrixd k11 (6, 6);
    for (std::size_t i{0}; i < 6; ++i)
    {
        k11 (i, i) = property.m_K[i];
    }
    
    kl.BlockFill (0, 0, k11);
    if (2 == nEleNodes)
    {

        kl.BlockFill (6, 6, k11);
        k11 = k11 * (-1.0);
        kl.BlockFill (0, 6, k11);
        kl.BlockFill (6, 0, k11);
    }
}

void BushCalculator::ComputeMass
(
const std::size_t                                    nEleNodes,
const SG::DataStructure::FEM::CBUSHData&             element,
const std::vector<SG::DataStructure::FEM::NodeData>& nodeSet,
const SG::DataStructure::FEM::PBUSHData&             property,
_OUT SG::Algebra::Matrixd&                           lumpedM
)
{
    lumpedM = Matrixd (nEleNodes * 6, nEleNodes * 6);
    if (1 == nEleNodes)
    {
        lumpedM (0, 0) = lumpedM (1, 1) = lumpedM (2, 2) = property.m_M;
    }
    else
    {
        Real alpha{element.m_s};
        if (element.m_oCId > -1)
        {
            const auto& coordProcessor = CoordinateProcessor::GetInstance ();

            // 计算节点在 Basic 坐标系下的坐标
            PointArray pointSet (2);
            std::vector<std::size_t> nodeIndexSet (nEleNodes);
            for (std::size_t iNode{0}; iNode < nEleNodes; ++iNode)
            {
                Id dispCoordId;
                coordProcessor.GetGlobalCoordinate (nodeSet[iNode], _OUT pointSet[iNode], _OUT dispCoordId);
            }

            // 计算偏置点在 Basic 坐标系下的坐标
            auto pointOffset = element.GetSi ();
            const Id offsetCoordId{element.m_oCId};
            if (offsetCoordId > 0)
            {
                coordProcessor.Translate2Basic (offsetCoordId, pointOffset, _OUT pointOffset);
            }

            // 计算质量分配系数 alpha
            const auto rA = length (pointSet[0], pointOffset);
            const auto rB = length (pointSet[1], pointOffset);
            alpha = rA / (rA + rB);
        }
        lumpedM (0, 0) = lumpedM (1, 1) = lumpedM (2, 2) = property.m_M * (1.0 - alpha);
        lumpedM (6, 6) = lumpedM (7, 7) = lumpedM (8, 8) = property.m_M * alpha;
    }
}

void BushCalculator::ComputeElementMatries (const SG::DBManager::DBServiceFactorySP dbServiceFactory)
{
    auto& assembly = AssemblyEigen::GetInstance ();

    CBUSHData element;
    std::size_t nEleNodes;
    std::size_t orienNodeIndex;
    std::vector<std::size_t> nodeIndexSet;
    std::vector<int>         dofGlobalIndexSet;
    std::vector<NodeData>    nodeSet (2);
    NodeData orienNode;
    static const DofSet dofSet{DOFDirections::DOF_Ux,   DOFDirections::DOF_Uy,   DOFDirections::DOF_Uz,
                               DOFDirections::DOF_Rotx, DOFDirections::DOF_Roty, DOFDirections::DOF_Rotz};

    const Real* K;

    bool isSpecifyEleCoordSys{true};

    const auto& eleService      = dbServiceFactory->get<IElementService> ();
    const auto  eleSize         = eleService->GetElementSize (m_type);
    const auto& propertyService = dbServiceFactory->get<IPBUSHService> ();
    const auto& nodeService     = dbServiceFactory->get<INodeService> ();

    for (std::size_t iEle{0}; iEle < eleSize; ++iEle)
    {
        // 获取单元数据
        eleService->FindByIndex (iEle, _OUT element);

        // 获取单元属性
        PBUSHData property;
        if (Status::STATUS_Failure == propertyService->FindById (element.m_pId, _OUT property))
        {
            std::cout << "Can't find Bush element property! \n";
            return;
        }
        K = property.m_K;

        // 判断单元节点数
        if (element.m_g[1] > 0)
        {
            nEleNodes = 2;
        }
        else if (0 == element.m_g[1])
        {
            nEleNodes = 1;
        }

        // 使用朝向节点时，获取朝向节点对应的索引，以及朝向节点数据
        if (element.m_go > 0)
        {
            assembly.FindNodeIndex (element.m_go, _OUT orienNodeIndex);
            nodeService->FindByIndex (orienNodeIndex, _OUT orienNode);
        }

        // 获取单元节点对应的索引，以及单元节点数据
        nodeIndexSet.resize (nEleNodes);
        for (std::size_t i{0}; i < nEleNodes; ++i)
        {
            assembly.FindNodeIndex (element.m_g[i], _OUT nodeIndexSet[i]);
            nodeService->FindByIndex (nodeIndexSet[i], _OUT nodeSet[i]);
        }

        // 获取单元自由度对应的全局自由度索引
        const std::size_t nDof{6 * nEleNodes};
        dofGlobalIndexSet.resize (nDof);
        assembly.GetGlobalDOFId (nodeIndexSet, dofSet, _OUT dofGlobalIndexSet);

        // 单元变换矩阵
        Matrixd TEle (nDof, nDof);
        Matrixd TEleOffset (nDof, nDof);
        ComputeElementTrans (nEleNodes, element, nodeSet, orienNode, K, _OUT isSpecifyEleCoordSys, _OUT TEle, _OUT TEleOffset);

        // 计算单元坐标系下单元刚度矩阵和质量矩阵
        Matrixd kg (nDof, nDof);
        ComputeStiffness (nEleNodes, property, _OUT kg);
        Matrixd lumpedM (nDof, nDof);
        ComputeMass (nEleNodes, element, nodeSet, property, lumpedM);

        // 未指定单元坐标系时，单元刚度矩阵仅需要单元坐标变换，注意单元质量矩阵也需要进行单元坐标变换
        if (isSpecifyEleCoordSys)
        {
            kg = TEleOffset.Transpose () * kg * TEleOffset;
        }
        else
        {
            kg = TEle.Transpose () * kg * TEle;
            lumpedM = TEle.Transpose () * lumpedM * TEle;
        }

        assembly.AccumulateStiffness (dofGlobalIndexSet, kg);
    }
}

void BushCalculator::ComputeElementInternalForce
(
const std::size_t           nEleNodes,
const SG::Algebra::Matrixd& kl,
const SG::Algebra::Matrixd& TEle,
const SG::Algebra::Matrixd& elementDisplacement,
_OUT SG::Algebra::Matrixd&  elementInternalForce
)
{
    // 单元坐标系单元内力
    const auto& internalForce = kl * TEle * elementDisplacement;

    if (1 == nEleNodes)
    {
        elementInternalForce = internalForce;
    }
    else
    {
        elementInternalForce = Matrixd (6, 1);
        for (std::size_t i{0}; i < 6; ++i)
        {
            {
                elementInternalForce (i, 0) = internalForce (6 + i, 0);
            }
        }
    }
}

void BushCalculator::ComputeElementStrain
(
const SG::DataStructure::FEM::PBUSHData& property,
const std::size_t                        nEleNodes,
const SG::Algebra::Matrixd&              TEle,
const SG::Algebra::Matrixd&              elementDisplacement,
_OUT SG::Algebra::Matrixd&               elementStrain
)
{
    // 计算单元坐标系位移
    const auto& eleDisplacementE = TEle * elementDisplacement;
    
    // 计算位移变化
    Matrixd deltaDisplacement (6, 1);
    if (1 == nEleNodes)
    {
        deltaDisplacement = eleDisplacementE;
    }
    else
    {
        for (std::size_t i{0}; i < 6; ++i)
        {
            deltaDisplacement (i, 0) = eleDisplacementE (6 + i, 0) - eleDisplacementE (i, 0);
        }
    }

    elementStrain = Matrixd (6, 1);
    for (std::size_t i{0}; i < 3; ++i)
    {
        elementStrain (i, 0) = deltaDisplacement (i, 0) * property.m_EA;
        elementStrain (3 + i, 0) = deltaDisplacement (3 + i, 0) * property.m_ET;
    }
}

void BushCalculator::ComputeElementStress
(
const SG::Algebra::Matrixd&              elementInternalForce,
const SG::DataStructure::FEM::PBUSHData& property,
_OUT SG::Algebra::Matrixd&               elementStress
)
{
    elementStress = Matrixd (6, 1);
    for (std::size_t i{0}; i < 3; ++i)
    {
        elementStress (i, 0) = elementInternalForce (i, 0) * property.m_SA;
        elementStress (3 + i, 0) = elementInternalForce (3 + i, 0) * property.m_ST;
    }
}

// 后处理数据
void BushCalculator::DataRecovery (const SG::DBManager::DBServiceFactorySP dbServiceFactory, int subcaseId, SG::DataStructure::Common::Real time, const std::vector<SG::DataStructure::Common::Real>& globalDisplacement)
{
    auto& assembly = AssemblyEigen::GetInstance ();

    CBUSHData element;
    std::size_t nEleNodes;
    std::size_t orienNodeIndex;
    std::vector<std::size_t> nodeIndexSet;
    std::vector<int>         dofGlobalIndexSet;
    std::vector<NodeData>    nodeSet (2);
    NodeData orienNode;
    static const DofSet dofSet{DOFDirections::DOF_Ux,   DOFDirections::DOF_Uy,   DOFDirections::DOF_Uz,
                               DOFDirections::DOF_Rotx, DOFDirections::DOF_Roty, DOFDirections::DOF_Rotz};

    const Real* K;

    bool isSpecifyEleCoordSys{true};
    Matrixd eleInternalForce (6, 1);
    Matrixd eleStrain (6, 1);
    Matrixd eleStress (6, 1);

    const auto& eleService      = dbServiceFactory->get<IElementService> ();
    const auto  eleSize         = eleService->GetElementSize (m_type);
    const auto& propertyService = dbServiceFactory->get<IPBUSHService> ();
    const auto& nodeService     = dbServiceFactory->get<INodeService> ();
    BUSHForceData eleForceData;
    std::vector<BUSHForceData> eleForceDataSet (eleSize);
    BUSHStrainData eleStrainData;
    std::vector<BUSHStrainData> eleStrainDataSet (eleSize);
    BUSHStressData eleStressData;
    std::vector<BUSHStressData> eleStressDataSet (eleSize);

    for (std::size_t iEle{0}; iEle < eleSize; ++iEle)
    {
        // 获取单元数据
        eleService->FindByIndex (iEle, _OUT element);

        // 获取单元属性
        PBUSHData property;
        if (Status::STATUS_Failure == propertyService->FindById (element.m_pId, _OUT property))
        {
            std::cout << "Can't find Bush element property! \n";
            return;
        }
        K = property.m_K;

        // 判断单元节点数
        if (element.m_g[1] > 0)
        {
            nEleNodes = 2;
        }
        else if (0 == element.m_g[1])
        {
            nEleNodes = 1;
        }

        // 使用朝向节点时，获取朝向节点对应的索引，以及朝向节点数据
        if (element.m_go > 0)
        {
            assembly.FindNodeIndex (element.m_go, _OUT orienNodeIndex);
            nodeService->FindByIndex (orienNodeIndex, _OUT orienNode);
        }

        // 获取单元节点对应的索引，以及单元节点数据
        nodeIndexSet.resize (nEleNodes);
        for (std::size_t i{0}; i < nEleNodes; ++i)
        {
            assembly.FindNodeIndex (element.m_g[i], _OUT nodeIndexSet[i]);
            nodeService->FindByIndex (nodeIndexSet[i], _OUT nodeSet[i]);
        }

        // 获取单元自由度对应的全局自由度索引
        const std::size_t nDof{6 * nEleNodes};
        dofGlobalIndexSet.resize (nDof);
        assembly.GetGlobalDOFId (nodeIndexSet, dofSet, _OUT dofGlobalIndexSet);

        // 单元变换矩阵：单元坐标系到位移输出坐标系变换 + 偏置变换
        Matrixd TEle (nDof, nDof);
        Matrixd TEleOffset (nDof, nDof);
        ComputeElementTrans (nEleNodes, element, nodeSet, orienNode, K, _OUT isSpecifyEleCoordSys, _OUT TEle, _OUT TEleOffset);

        // 计算单元坐标系下单元刚度矩阵
        Matrixd kl (nDof, nDof);
        ComputeStiffness (nEleNodes, property, _OUT kl);

        // 从总体位移向量中提取单元位移
        Matrixd eleDisplacement (nDof, 1);
        for (std::size_t iDof{0}; iDof < nDof; ++iDof)
        {
            eleDisplacement (iDof, 0) = globalDisplacement[dofGlobalIndexSet[iDof]];
        }

        // 计算单元内力
        ComputeElementInternalForce (nEleNodes, kl, TEleOffset, eleDisplacement, _OUT eleInternalForce);

        // 转换为存储的单元内力数据结构
        eleForceData.m_eId = element.m_id;
        eleForceData.m_FX  = eleInternalForce (0, 0);
        eleForceData.m_FY  = eleInternalForce (1, 0);
        eleForceData.m_FZ  = eleInternalForce (2, 0);
        eleForceData.m_MX  = eleInternalForce (3, 0);
        eleForceData.m_MY  = eleInternalForce (4, 0);
        eleForceData.m_MZ  = eleInternalForce (5, 0);
        eleForceDataSet[iEle] = eleForceData;

        // 计算单元应变
        ComputeElementStrain (property, nEleNodes, TEleOffset, eleDisplacement, _OUT eleStrain);

        // 转换为存储的单元应变数据结构
        eleStrainData.m_eId = element.m_id;
        eleStrainData.m_TX  = eleStrain (0, 0);
        eleStrainData.m_TY  = eleStrain (1, 0);
        eleStrainData.m_TZ  = eleStrain (2, 0);
        eleStrainData.m_RX  = eleStrain (3, 0);
        eleStrainData.m_RY  = eleStrain (4, 0);
        eleStrainData.m_RZ  = eleStrain (5, 0);
        eleStrainDataSet[iEle] = eleStrainData;

        // 计算单元应力
        ComputeElementStress (eleInternalForce, property, _OUT eleStress);

        // 转换为存储的单元应力数据结构
        eleStressData.m_eId = element.m_id;
        eleStressData.m_TX  = eleStress (0, 0);
        eleStressData.m_TY  = eleStress (1, 0);
        eleStressData.m_TZ  = eleStress (2, 0);
        eleStressData.m_RX  = eleStress (3, 0);
        eleStressData.m_RY  = eleStress (4, 0);
        eleStressData.m_RZ  = eleStress (5, 0);
        eleStressDataSet[iEle] = eleStressData;
    }

    dbServiceFactory->get<IBUSHForceService> ()->Add (SG::DBManager::Post::DomainRequest (subcaseId), eleForceDataSet);
    dbServiceFactory->get<IBUSHStrainService> ()->Add (SG::DBManager::Post::DomainRequest (subcaseId), eleStrainDataSet);
    dbServiceFactory->get<IBUSHStressService> ()->Add (SG::DBManager::Post::DomainRequest (subcaseId), eleStressDataSet);
}

// 计算位移输出坐标系下应力刚度矩阵
void BushCalculator::ComputeStressStiffness (const SG::DBManager::DBServiceFactorySP dbServiceFactory, int staticSubId)
{
    
}